from os import getcwd, path
from numpy import linspace, abs, zeros, real, arange, angle
from numpy.fft import fft, fftfreq, fftshift
from enthought.traits.api import Any, Bool, Int, Float, Str, Range, Button, HasTraits, \
                                DelegatesTo, Instance, Enum, Event, Trait, Tuple, \
                                Callable, Array, List, Property, on_trait_change
from enthought.traits.ui.api import View, Item, Group, HGroup, VGroup, VGrid, \
                                    RangeEditor, HSplit, VFold, HFlow
from enthought.traits.ui.menu import OKButton, CancelButton
from enthought.enable.api import BaseTool, KeySpec, ColorTrait, LineStyle, \
                                 Component, ComponentEditor, Window
from enthought.chaco.api import Plot, ArrayPlotData, DataRange1D, \
                                ArrayDataSource, LinearMapper, BasePlotContainer, \
                                OverlayPlotContainer, VPlotContainer, LinePlot, DataLabel, \
                                create_line_plot, create_scatter_plot, PlotLabel, \
                                add_default_axes
from enthought.chaco.tools.api import PanTool, SimpleZoom
from mpl_figure_editor import MPLFigureEditor, Figure
from pylab import figure, plot, xlabel, ylabel

from gprViewerUIs import wxOpenFile


class LoadPanel(HasTraits):
    file = Str()
    dir = Str()
    browse = Button(label='...')
    sample_min = Int(0, label="Sample min/max")
    sample_max = Int(-1)
    trace_no = Int(0)
    trace_max = Int(-1, label="Max:")
    show_button = Button(label="Show")
    show = Bool(False)
    
    view = View(VGrid(
                       Item(name='file'), Item(name='browse', show_label=False),
                       Item(name='sample_min'), Item(name='sample_max', show_label=False),
                       Item(name='trace_no'), Item(name='trace_max', style='readonly'),
                       Item(name='show_button', show_label=False),
                       show_border=True, label="Choose data file and range"),
                       width=100
                       )
    

class TraceViewer(HasTraits):
    # Load Panel Items
    figure = Instance(Figure, ())
#    file = Str()
#    dir = Str()
#    browse = Button(label='...')
#    sample_min = Int(0, label="Sample min/max")
#    sample_max = Int(-1)
#    trace_no = Int(0)
#    trace_max = Int(-1, label="Max:")
#    show_button = Button(label="Show")
#    show = Bool(False)
    # Processing Panel Items
    spectrum_button = Button(label="Spectrum")
    cspectrum_button = Button(label="Complex Spectrum")
    # Status Panel
    status_txt = Str("TraceViewer started ...")
    # Plot Panel Items
    container = Instance(BasePlotContainer)
    
    load_panel = Instance(LoadPanel, ())
    
#    load_panel = VGrid(
#                       Item(name='file'), Item(name='browse', show_label=False),
#                       Item(name='sample_min'), Item(name='sample_max', show_label=False),
#                       Item(name='trace_no'), Item(name='trace_max', style='readonly'),
#                       Item(name='show_button', show_label=False),
#                       show_border=True, label="Choose data file and range")
    
    proc_panel = VGroup(
                        Item(name='spectrum_button', show_label=False, width=300),
                        Item(name='cspectrum_button', show_label=False),
                        show_border=True
                        )
    
    status_panel = VGroup(
                          Item(name='status_txt', show_label=False, style='custom',
                               full_size=False, height=100)
                          )
    
    view = View(HSplit(
                       VGroup(
                              Item(name="load_panel", style="custom", show_label=False), proc_panel, status_panel
#                              load_panel, proc_panel, status_panel
                              ),
#                      '_',
#                      Item(name="container", show_label=False, editor=ComponentEditor())
#                          visible_when="show==True")),
                      Item('figure', editor=MPLFigureEditor(), show_label=False),
                       ),
                icon=r"../src/recources/icons/edit-32.png",
                title="Trace Viewer",
                width=800, height=0.9,
                resizable=True)
    
    def __init__(self, file="", dir=getcwd()):
        self.dir = dir
        self.file = file
        
    def _create_traceplot(self):
        x = self.x
        y = self.y
        print x.shape, y.shape
        fig = figure()
        plot(x,y)
        
#        axes = self.figure.add_subplot(111)
#        self.axes.plot(x,y)
        xlabel("Time [ns]")
        ylabel("Signal amplitude")
        self.figure = fig
        self.figure.canvas.draw()        
        
    def _create_spectrumplot(self, x, y):
        
        axes.plot(x,y)
#        axes.xlabel("Frequency [MHz]")
#        axes.ylabel("Amplitude")
        self.figure.canvas.draw()
        
    def _create_complexspectrumplot(self,x, y1, y2):
                
        amplitudeplot = create_line_plot((x, y1), index_bounds = None, value_bounds = None,
                         orientation = 'h', color = 'green', width = 1.0, dash = 'solid',
                         value_mapper_class = LinearMapper,
                         bgcolor = 'white', border_visible = True,
                         add_grid = False, add_axis = False, index_sort = 'ascending')
        add_default_axes(amplitudeplot, orientation = 'normal', htitle = 'Frequency [MHz]', vtitle = 'Amplitude')
#        amplitudeplot.origin = "top left"
        amplitudeplot.tools.append(PanTool(amplitudeplot, drag_button="right"))
        zoom = SimpleZoom(component=amplitudeplot, tool_mode="box", drag_button="left", always_on=True)
        amplitudeplot.overlays.append(zoom)
        
        phaseplot = create_line_plot((x, y2), index_bounds = None, value_bounds = None,
                         orientation = 'h', color = 'red', width = 1.0, dash = 'solid',
                         value_mapper_class = LinearMapper,
                         bgcolor = 'white', border_visible = True,
                         add_grid = False, add_axis = False, index_sort = 'ascending')
        add_default_axes(phaseplot, orientation = 'normal', htitle = 'Frequency [MHz]', vtitle = 'Phase')
#        phaseplot.origin = "top left"
        phaseplot.tools.append(PanTool(phaseplot, drag_button="right"))
        zoom = SimpleZoom(component=phaseplot, tool_mode="box", drag_button="left", always_on=True)
        phaseplot.overlays.append(zoom)
        
        container = VPlotContainer(padding=40, padding_left=60, spacing=40)
        container.add(phaseplot)
        container.add(amplitudeplot)
        self.container = container
        
    def _file_changed(self):
        try:
            from tables import openFile, NoSuchNodeError
            h5file = openFile(path.join(self.dir,self.file), mode = 'r')
            try:
                sample_max, trace_max = h5file.root.data.traces.shape
                self.sample_max = int(sample_max)
                self.trace_max = int(trace_max)
            except NoSuchNodeError:
                self.sample_max, self.trace_max = h5file.root.profile.traces.shape
            h5file.close()
        except IOError:
            print "The chosen file has an unknown format."
    
    def _browse_fired(self):
        self.file, self.dir = wxOpenFile(self.dir, multi=False)
        
    def _show_button_fired(self):
        from processGPR import HUBRA
#        self.show = not self.show
        gpr = HUBRA()
        gpr.readH5(path.join(self.dir,self.file),
                   self.trace_no, self.trace_no,
                   self.sample_min, self.sample_max)
        self.x = gpr.t
        self.y = gpr.data
        self.y.shape = (gpr.data.shape[0],)
        self.deltaT = gpr.deltaT
        self._create_traceplot()
        self.status_txt = "%s\nTrace %i from file %s displayed" % (self.status_txt, self.trace_no, self.file) 
        
    def _spectrum_button_fired(self):
        Nfft = self.y.size
        freq = fftfreq(Nfft, self.deltaT)
        Y = abs(fft(self.y))
        x = freq[0:Nfft/2]*1e-6
        y = Y[0:Nfft/2]
        self._create_spectrumplot(x, y)
        
    def _cspectrum_button_fired(self):
        Nfft = self.y.size
        freq = fftfreq(Nfft, self.deltaT)
        Y = fft(self.y)
        x = freq[0:Nfft/2]*1e-6
        y1 = abs(Y)[0:Nfft/2]
        y2 = angle(Y)[0:Nfft/2]
        self._create_complexspectrumplot(x, y1, y2)
        
if __name__ == "__main__":
    file = '20090212_stk01_1000ns_phase.h5'
    dir = r'Y:\Data2008_09\Hubra_proc\IPY'
    tp = TraceViewer(file, dir)
    tp.configure_traits()
    